import os
import time
import ctypes
import numpy as np
from typing import List

cur_path = os.path.split(__file__)[0]
lib_path = os.environ.get("ACL_LIB_PATH", f"{cur_path}/build/libutils.so")
libutils = ctypes.CDLL(lib_path)

# env
libutils.creatEnv.argtypes = [ctypes.c_int32]
libutils.creatEnv.restype = ctypes.c_void_p
libutils.destroyEnv.argtypes = [ctypes.c_void_p]
libutils.destroyEnv.restype = None
libutils.envSyncStream.argtypes = [ctypes.c_void_p]
libutils.envSyncStream.restype = ctypes.c_int

# NDTensor
libutils.createNDTensor.argtypes = [
    ctypes.c_void_p,
    ctypes.c_uint64,
    ctypes.c_int,
    ctypes.c_void_p,
]
libutils.createNDTensor.restype = ctypes.c_void_p
libutils.destroyNDTensor.argtypes = [ctypes.c_void_p]
libutils.destroyNDTensor.restype = None
libutils.ndTensor2Cpu.argtypes = [ctypes.c_void_p]
libutils.ndTensor2Cpu.restype = None


class AclEnv:
    def __init__(self, device_id=0):
        self.ptr = libutils.creatEnv(device_id)

    def __del__(self):
        libutils.destroyEnv(self.ptr)

    def sync(self):
        libutils.envSyncStream(self.ptr)


class AclNDTensor:
    def __init__(self, np_array: np.ndarray):
        shape = np.array(np_array.shape, dtype=np.int64)
        shape_size = len(np_array.shape)
        dev_ptr = np_array.ctypes.data
        acl_dtype = numpy_dtype_2_acl_dtype(np_array.dtype)
        self.ptr = libutils.createNDTensor(
            shape.ctypes.data, shape_size, acl_dtype, dev_ptr
        )

    def __del__(self):
        libutils.destroyNDTensor(self.ptr)

    def to_cpu(self):
        libutils.ndTensor2Cpu(self.ptr)


class OpRunner:
    def __init__(
        self, op_path, tensor_count: int, env=AclEnv(0), pint_time=False
    ) -> None:
        self.op = ctypes.CDLL(op_path)
        self.op.run.argtypes = [ctypes.c_void_p for i in range(tensor_count + 1)]
        self.op.run.restype = ctypes.c_int
        self.fn = self.op.run
        self.env = env
        self.print_time = pint_time

    def __call__(self, tensors: List[np.ndarray], output_index: List[int]):
        acl_tensors = []
        for tensor in tensors:
            if isinstance(tensor, np.ndarray):
                acl_tensors.append(AclNDTensor(tensor))
                continue
            acl_tensors.append(tensor)

        input_tensors = []
        for t in acl_tensors:
            if isinstance(t, AclNDTensor):
                input_tensors.append(t.ptr)
                continue
            input_tensors.append(t)
        if self.print_time:
            pre = time.time_ns()
        self.fn(self.env.ptr, *input_tensors)
        self.env.sync()
        if self.print_time:
            print("运行Cost:", (time.time_ns() - pre) / 1000 / 1000, "毫秒(ms)")
        for idx in output_index:
            acl_tensors[idx].to_cpu()


def get_loss_by_type(dtype):
    loss = 0
    if dtype == np.float16:
        loss = 1 / 1000
    elif dtype == np.float32:
        loss = 1 / 10000
    return loss


def _compare(v1: np.ndarray, v2: np.ndarray):
    loss = get_loss_by_type(v1.dtype)
    return np.abs(v1 - v2) <= loss


def compare(v1: np.ndarray, v2: np.ndarray):
    return _compare(v1, v2).all()


def right_rate(v1: np.ndarray, v2: np.ndarray):
    ret = _compare(v1, v2)
    return ret.astype(np.int32).sum() / v1.size


def verify_result(real_result, golden):
    loss = get_loss_by_type(real_result.dtype)
    minimum = 10e-10
    result = np.abs(real_result - golden)  # 计算运算结果和预期结果偏差
    deno = np.maximum(np.abs(real_result), np.abs(golden))  # 获取最大值并组成新数组
    result_atol = np.less_equal(result, loss)  # 计算绝对误差
    result_rtol = np.less_equal(result / np.add(deno, minimum), loss)  # 计算相对误差
    if not result_rtol.all() and not result_atol.all():
        if (
            np.sum(result_rtol == False) > real_result.size * loss
            and np.sum(result_atol == False) > real_result.size * loss
        ):  # 误差超出预期时返回打印错误，返回对比失败
            print("[ERROR] result error")
            return False
    print("test pass")
    return result_atol, result_rtol


class AclDtype:
    ACL_FLOAT = 0
    ACL_FLOAT16 = 1
    ACL_INT8 = 2
    ACL_INT32 = 3
    ACL_UINT8 = 4
    ACL_INT16 = 6
    ACL_UINT16 = 7
    ACL_UINT32 = 8
    ACL_INT64 = 9
    ACL_UINT64 = 10
    ACL_DOUBLE = 11
    ACL_BOOL = 12
    ACL_STRING = 13
    ACL_COMPLEX64 = 16
    ACL_COMPLEX128 = 17
    ACL_BF16 = 27
    ACL_INT4 = 29
    ACL_UINT1 = 30
    ACL_COMPLEX32 = 33


class AclRunMode:
    ACL_DEVICE = 0
    ACL_HOST = 1


class AclAllocPolicy:
    """
    0：ACL_MEM_MALLOC_HUGE_FIRST，当申请的内存小于等于1M时，即使使用该内存分配规则，也是申请普通页的内存。当申请的内存大于1M时，优先申请大页内存，如果大页内存不够，则使用普通页的内存。
    1：ACL_MEM_MALLOC_HUGE_ONLY，仅申请大页，如果大页内存不够，则返回错误。
    2：ACL_MEM_MALLOC_NORMAL_ONLY，仅申请普通页。
    3：ACL_MEM_MALLOC_HUGE_FIRST_P2P，仅Device之间内存复制场景下申请内存时使用该选项，表示优先申请大页内存，如果大页内存不够，则使用普通页的内存。预留选项。
    4：ACL_MEM_MALLOC_HUGE_ONLY_P2P，仅Device之间内存复制场景下申请内存时使用该选项，仅申请大页内存，如果大页内存不够，则返回错误。预留选项。
    5：ACL_MEM_MALLOC_NORMAL_ONLY_P2P，仅Device之间内存复制场景下申请内存时使用该选项，仅申请普通页的内存。预留选项。
    """

    ACL_MEM_MALLOC_HUGE_FIRST = 0
    ACL_MEM_MALLOC_HUGE_ONLY = 1
    ACL_MEM_MALLOC_NORMAL_ONLY = 2
    ACL_MEM_MALLOC_HUGE_FIRST_P2P = 3
    ACL_MEM_MALLOC_HUGE_ONLY_P2P = 4
    ACL_MEM_MALLOC_NORMAL_ONLY_P2P = 5


class AclMemcopyKind:
    """
    0：ACL_MEMCPY_HOST_TO_HOST，Host内的内存复制。
    1：ACL_MEMCPY_HOST_TO_DEVICE，Host到Device的内存复制。
    2：ACL_MEMCPY_DEVICE_TO_HOST，Device到Host的内存复制。
    3：ACL_MEMCPY_DEVICE_TO_DEVICE，Device内的内存复制。
    """

    ACL_MEMCPY_HOST_TO_HOST = 0
    ACL_MEMCPY_HOST_TO_DEVICE = 1
    ACL_MEMCPY_DEVICE_TO_HOST = 2
    ACL_MEMCPY_DEVICE_TO_DEVICE = 3


def numpy_dtype_2_acl_dtype(numpy_dtype):
    if numpy_dtype == np.float32:
        return AclDtype.ACL_FLOAT
    if numpy_dtype == np.float16:
        return AclDtype.ACL_FLOAT16
    if numpy_dtype == np.int8:
        return AclDtype.ACL_INT8
    if numpy_dtype == np.int32:
        return AclDtype.ACL_INT32
    if numpy_dtype == np.uint8:
        return AclDtype.ACL_UINT8
    if numpy_dtype == np.int16:
        return AclDtype.ACL_INT16
    if numpy_dtype == np.uint16:
        return AclDtype.ACL_UINT16
    if numpy_dtype == np.uint32:
        return AclDtype.ACL_UINT32
    if numpy_dtype == np.int64:
        return AclDtype.ACL_INT64
    if numpy_dtype == np.uint64:
        return AclDtype.ACL_UINT64
    if numpy_dtype == np.double:
        return AclDtype.ACL_DOUBLE
    if numpy_dtype == np.bool_:
        return AclDtype.ACL_BOOL
    if numpy_dtype == np.string_:
        return AclDtype.ACL_STRING
    if numpy_dtype == np.complex64:
        return AclDtype.ACL_COMPLEX64
    if numpy_dtype == np.complex128:
        return AclDtype.ACL_COMPLEX128
    if numpy_dtype == np.complex_:
        return AclDtype.ACL_COMPLEX32
    # TODO 如何使用bf16


if __name__ == "__main__":
    a = np.zeros((3, 3), dtype=np.float16)
    env = AclEnv(0)
    nd_tensor = AclNDTensor(a)
